#ifndef POWERSPECTRUM_H
#define POWERSPECTRUM_H

#include "MoreFunctional.h"

#include "Save.h"
#include <map>
#include <iostream>

#include "dbg.h"

class Universe;

double rawg(const Universe& u, double w);
double rawpf(const Universe& u, double w);
double raww_z(const Universe& u, double z);
double rawdwda(const Universe& u, double a);
double rawpom(const Universe& u, double l);
double rawpgam(const Universe& u, double l);
double rawpkappa(const Universe& u, double l);
double rawfom(const Universe& u, double R);
double rawfgam(const Universe& u, double R);

double calc_hsq(double om_m, double om_v, double w, double wa, double a);
bool NegHsq(double om_m, double om_v, double w, double wa);

double WMap(double eta);
double WVar(double eta);

typedef binder2_1<std::pointer_to_binary_function<const Universe&,double,double> > BOUND_F;

struct UParams { 
    UParams(
        double _b, double _c, double _n, double _w, double _wa,
        double _A, double _h) :
        ombh2(_b), omch2(_c), n(_n), w(_w), wa(_wa), As(_A), h(_h)
    {}

    double ombh2, omch2, n, w, wa, As, h; // primary values;

    double GetOmM() const { return (omch2+ombh2)/h/h; }
    double GetOmB() const { return (ombh2)/h/h; }
    double GetOmDE() const { return 1.-GetOmM(); }

    bool operator==(const UParams& rhs) const
    {
        return (ombh2 == rhs.ombh2 && omch2 == rhs.omch2 && n == rhs.n && 
                w == rhs.w && wa == rhs.wa && As == rhs.As && h == rhs.h);
    }

    bool operator<(const UParams& rhs) const
    { 
        if (ombh2 == rhs.ombh2) {
            if (omch2 == rhs.omch2) {
                if (n == rhs.n) {
                    if (w == rhs.w) {
                        if (wa == rhs.wa) {
                            if (As == rhs.As) {
                                return h < rhs.h;
                            } else return (As < rhs.As);
                        } else return (wa < rhs.wa);
                    } else return (w < rhs.w);
                } else return (n < rhs.n);
            } else return (omch2 < rhs.omch2);
        } else return (ombh2 < rhs.ombh2);
    }

};

struct UParamsNonFlat { 
    UParamsNonFlat(
        double _b, double _c, double _de, double _n, double _w, double _wa, 
        double _A, double _h) :
        ombh2(_b), omch2(_c), omde(_de), n(_n), w(_w), wa(_wa), As(_A), h(_h)
    {}

    double ombh2, omch2, omde, n, w, wa, As, h; // primary values;

    double GetOmM() const { return (omch2+ombh2)/h/h; }
    double GetOmB() const { return (ombh2)/h/h; }
    double GetOmDE() const { return omde; }

    bool operator==(const UParamsNonFlat& rhs) const
    {
        return (ombh2 == rhs.ombh2 && omch2 == rhs.omch2 && n == rhs.n && 
                w == rhs.w && wa == rhs.wa && As == rhs.As && h == rhs.h);
    }

    bool operator<(const UParamsNonFlat& rhs) const
    { 
        if (ombh2 == rhs.ombh2) {
            if (omch2 == rhs.omch2) {
                if (omde == rhs.omde) {
                    if (n == rhs.n) {
                        if (w == rhs.w) {
                            if (wa == rhs.wa) {
                                if (As == rhs.As) {
                                    return h < rhs.h;
                                } else return (As < rhs.As);
                            } else return (wa < rhs.wa);
                        } else return (w < rhs.w);
                    } else return (n < rhs.n);
                } else return (omde < rhs.omde);
            } else return (omch2 < rhs.omch2);
        } else return (ombh2 < rhs.ombh2);
    }

};

class Universe {

public:

    Universe(
        double om_m, double om_v, double om_b, double h0,
        double As, double n, double w, double wa);
    Universe(const UParams& up);
    Universe(const UParamsNonFlat& up);
    ~Universe() {}

    double NLP3d(double k, double z) const;
    double NLDelta3d(double k, double z) const;
    // 3D power spectrum using the Smith et al formulation as modified
    // by Derek to handle w, wa
    // Delta = P(k,z) * k^3 / 2Pi^2
    double LinearP3d(double k, double z) const;
    double LinearDelta3d(double k, double z) const;
    // The linear 3D power spectrum

    double fk(double w) const;
    // fk(w) is the comoving angular diameter distance at comoving distance w
    //
    // fk(w) = w if K = 0
    //       =  sqrt(K)^-1 sin( sqrt(K) w) if K > 0
    //       = sqrt(-K)^-1 sin(sqrt(-K) w) if K < 0
    //
    // Note: we use the usual dimensionless version of w (= w*H0/c)
    // so we need to multiply w by coverH, then divide fk(w) by the same
    // This is equivalent to using sqrt(K)*coverH where sqrt(K) appears above
    //

    double w_z(double z) const { return sw_z(z); }
    // coordinate distance (in units of c/Ho) as a function of redshift
    // w = int( 1/sqrt(Om_m(1+z)^3 + Om_k(1+z)^2 + Om_v) , z=0..z0 )

    double z_w(double w) const;
    // redshift as a function of coordinate distance (in units of c/Ho)

    double dwdz(double z) const
        // = dwda dadz = dwda (-(1+z)^-2) = -dwda * a^2
        // dwda actually returns the negative of the real value, so omit - here.
    { 
        double a = 1./(1.+z);
        return rawdwda(*this,a)*(a*a);
    }

    double hsq(double a) const
    { return calc_hsq(om_m,om_v,w,wa,a); }

    double g(double w) const;
    // lensing strength of lens at distance w
    
    double pf(double w) const;
    // smoothed estimate of foreground galaxy density 
    // pfsmooth = int_w-dw,w+dw pf(w) dw
    
    double pb(double w) const;
    // smoothed estimate of background galaxy density 
    // pfsmooth = int_w-dw,w+dw pf(w) dw

    double fnn(double R) const;
    // <N^2> = fnn b^2
    // fnn(R) = int dl l pom(l) W(lR)

    double fnm(double R) const;
    // <NM> = fnm br
    // fnm(R) = int dl l pgam(l) W(lR)

    double fmm(double R) const;
    // <M^2> = fmm 
    // fmm(R) = int dl l pkappa(l) W(lR)

    double fvar(double R) const;
    // <gamma^2> = fvar
    // fvar(R) = int dl l pkappa(l) W_TH(lR)

    double pom(double l) const { return spom(l); }

    double fom(double R) const { return sfom(R); }
    // omega(R) = fom b^2
    // fom(R) = 1/2Pi int dl l J0(lR) int dw pf(w)^2/fk(w)^2 P3D(l/fk(w);w)

    double pgam(double l) const { return spgam(l); }
    // pgam(l) = int dw pf(w)g(w)/a(w)fk(w) P3D(l/fk(w);w)

    double fgam(double R) const { return sfgam(R); }
    // <gamma_t>(R) = fgam b^2
    // fgam(R) = 1/2Pi int dl l J0(lR) pgam(l)

    double pkappa(double l) const { return spkappa(l); }
    // pkappa(l) = int dw g(w)^2/a(w)^2 P3D(l/fk(w);w)

    double altfnm(double R) const;
    // fnm = int dr r/R^2 fgam(r) Tcross(r/R)

    double altfnn(double R) const;
    // fnn = int dr r/R^2 fom(r) Tcross(r/R)

    double sigma8() const { return sigma(8.); }
    double sigma(double r) const;
    // sigma8 is the usual sigma_8 parameter
    // sigma(r) is for other scales, measured in h^-1 Mpc

    const double om_m;
    const double om_v;
    const double om_b;
    const double h;
    const double A_s;
    const double n;
    const double w;
    const double wa;
    const double coverH;
    const double coverHcb;
    double K;
    bool flat;
    double sqrtK;
    mutable double pfint;
    mutable double pbint;

    mutable std::multimap<double,double> pbmap;
    mutable std::multimap<double,double> pfmap;
    // These maps are such that each entry is a pair
    // of w (the comoving distance in dimensionless form - no c/H)
    // and weight.  The weight is based on the relative number of
    // galaxies in the redshift survey at each magnitude compared
    // to the total weight in observed galaxies in our lensing
    // survey at the same magnitudes.
    // 
    // For practical purposes, if pb(w) is the real underlying
    // density of galaxies as a function of w, and Wi is the
    // returned weights in the multimap, then
    // int dw pb(w) --> sum_i Wi

    mutable std::map<double,double> zwmap;

    bool precise;
    mutable bool def;

private:

    SavedFunctionLinear<double,BOUND_F> sg;
    SavedFunctionLinear<double,BOUND_F> spf;
    SavedFunctionLinear<double,BOUND_F> spb;
    SavedFunctionLinear<double,BOUND_F> sw_z;
    SavedFunctionLogLog<double,BOUND_F> spom;
    SavedFunctionLogLog<double,BOUND_F> spgam;
    SavedFunctionLogLog<double,BOUND_F> spkappa;
    SavedFunctionLogLog<double,BOUND_F> sfom;
    SavedFunctionLogLog<double,BOUND_F> sfgam;

private:
    Universe(const Universe& rhs);
    void operator=(const Universe& rhs);

};


#endif
